package com.aiwiown.snackmq.storage.partition;

import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.storage.segment.LogSegment;
import com.aiwiown.snackmq.storage.segment.SegmentFullException;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentNavigableMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * 单个分区的存储实现。
 * 作为一个“分区经理”，它管理着多个 LogSegment（日志段）。
 */
@Slf4j
public class PartitionStorage implements AutoCloseable {

    private final String partitionPath;
    private final long segmentSize;
    private final int indexInterval;
    // 【新增】存储恢复策略配置
    private final boolean forceFullLogRecovery;
    // 这个 map 持有所有的“日志工人”（LogSegment），按基准偏移量排序
    private final ConcurrentNavigableMap<Long, LogSegment> segments = new ConcurrentSkipListMap<>();
    // 这是整个分区的下一个逻辑偏移量
    private final AtomicLong nextOffset = new AtomicLong(0);
    // 一个读写锁，用于在创建新段时保护 segments 集合
    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();

    public PartitionStorage(String basePath, String topic, int partitionId, long segmentSize, int indexInterval, boolean forceFullLogRecovery) {
        this.partitionPath = basePath + File.separator + topic + "-" + partitionId;
        this.segmentSize = segmentSize;
        this.indexInterval = indexInterval;
        // 【新增】赋值
        this.forceFullLogRecovery = forceFullLogRecovery;
    }

    /**
     * 加载分区，从磁盘恢复现有的日志段。
     */
    public void load() throws IOException {
        File dir = new File(partitionPath);
        if (!dir.exists() && !dir.mkdirs()) {
            throw new IOException("无法创建分区目录: " + partitionPath);
        }

        File[] files = dir.listFiles((d, name) -> name.endsWith(".log"));
        if (files == null || files.length == 0) {
            // 如果是新分区，创建第一个日志段
            createNewSegment(0);
        } else {
            for (File file : files) {
                long baseOffset = Long.parseLong(file.getName().replace(".log", ""));
                LogSegment segment = new LogSegment(partitionPath, baseOffset, segmentSize, indexInterval, this.forceFullLogRecovery);
                segments.put(baseOffset, segment);
            }
            // 根据最后一个段的状态，设置下一个逻辑偏移量
            LogSegment lastSegment = segments.lastEntry().getValue();
            nextOffset.set(lastSegment.getBaseOffset() + lastSegment.getEntryCount().get());
        }
        log.info("分区 {} 加载完成。下一个偏移量是: {}", partitionPath, nextOffset.get());
    }

    /**
     * 创建一个新的日志段。这是一个写操作，需要获取写锁。
     */
    private LogSegment createNewSegment(long baseOffset) throws IOException {
        lock.writeLock().lock();
        try {
            LogSegment newSegment = new LogSegment(partitionPath, baseOffset, segmentSize, indexInterval, this.forceFullLogRecovery);
            segments.put(baseOffset, newSegment);
            log.info("为分区 {} 创建了新的日志段，基准偏移量: {}", partitionPath, baseOffset);
            return newSegment;
        } finally {
            lock.writeLock().unlock();
        }
    }

    /**
     * Leader 节点使用的追加方法。
     */
    public long append(Message message) throws IOException {
        long logicalOffset = nextOffset.get();
        append(message, logicalOffset);
        return logicalOffset;
    }

    /**
     * Follower 节点在复制数据时，或被 Leader 节点内部调用的追加方法。
     */
    public void append(Message message, long logicalOffset) throws IOException {
        lock.readLock().lock();
        try {
            // 确保我们是按顺序追加的
            if (logicalOffset != this.nextOffset.get()) {
                log.warn("在分区 {} 检测到乱序追加。期望偏移量: {}, 实际偏移量: {}. 这在副本追赶数据时可能发生，将重置偏移量。",
                        this.partitionPath, this.nextOffset.get(), logicalOffset);
                this.nextOffset.set(logicalOffset);
            }

            // 获取当前活跃的“工人”（最后一个日志段）
            LogSegment segment = segments.lastEntry().getValue();
            try {
                // 将工作委托给“工人”
                segment.append(message, logicalOffset);
                this.nextOffset.incrementAndGet(); // 成功追加后，更新下一个偏移量
            } catch (SegmentFullException e) {
                // 当前“工人”累了（段已满），需要雇佣一个新的
                lock.readLock().unlock(); // 释放读锁，准备获取写锁
                lock.writeLock().lock();
                try {
                    // 双重检查，防止其他线程已经创建了新段
                    if (segment == segments.lastEntry().getValue()) {
                        segment.flush();
                        segment = createNewSegment(logicalOffset);
                    } else {
                        segment = segments.lastEntry().getValue();
                    }
                    // 在新段上重试追加
                    segment.append(message, logicalOffset);
                    this.nextOffset.incrementAndGet();
                } finally {
                    lock.writeLock().unlock();
                    lock.readLock().lock(); // 锁降级，重新获取读锁
                }
            }
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * 从分区中读取一条消息。
     */
    public Message read(long offset) throws IOException {
        lock.readLock().lock();
        try {
            // 找到负责该偏移量的“工人”
            Map.Entry<Long, LogSegment> entry = segments.floorEntry(offset);
            if (entry != null) {
                // 将工作委托给“工人”
                return entry.getValue().read(offset);
            }
            return null;
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * 获取下一个将要写入的逻辑偏移量。
     */
    public long getNextOffset() {
        return this.nextOffset.get();
    }

    /**
     * 关闭整个分区，包括它管理的所有日志段。
     */
    @Override
    public void close() throws IOException {
        lock.writeLock().lock();
        try {
            log.info("正在关闭分区 {}", partitionPath);
            for (LogSegment segment : segments.values()) {
                segment.close();
            }
            segments.clear();
        } finally {
            lock.writeLock().unlock();
        }
    }
}